import torch
from torch.utils.data import DataLoader
from torch import nn, optim

from model.dnn.net import Net01
from model.dnn import dataset
import model.dnn.val as val

import warnings
warnings.filterwarnings("ignore")

from tools.plot_model import plot_singleModel


def nural_networks(filePath):
    accuracy_list, precision_list, recall_list, f1_list = [], [], [], []
    decive = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    net = Net01().to(decive)

    # 优化器的作用在于优化学习率
    opt = optim.Adam(net.parameters())

    # 损失函数
    bceloss = nn.BCELoss(reduction="mean")

    # 数据集
    train_datasets = dataset.NasaDataset(datafile=filePath,
                                         isTrain=True)
    test_datasets = dataset.NasaDataset(datafile=filePath,
                                        isTrain=False)

    # 加载器
    train_dataloader = DataLoader(train_datasets, batch_size=50, shuffle=True)
    test_dataloader = DataLoader(test_datasets, batch_size=50, shuffle=True)

    for epoch in range(1001):
        net.train()
        train_losses = []
        # 这里只有一个epoch
        for i, (data, target) in enumerate(train_dataloader):
            data, target = data.to(decive), target.to(decive)

            pred = net(data)

            loss = bceloss(pred, target)

            opt.zero_grad()  # 清空梯度
            loss.requires_grad_(True)  # 不加上这么一行代码没办法进行反向更新
            loss.backward()  # 反向求导
            opt.step()  # 更新梯度

            train_losses.append(loss)

        epoch_losses = torch.mean(torch.tensor(train_losses))
        # print("epoch {} train_losses = {}".format(epoch, epoch_losses))

        if epoch % 100 == 0:
            net.eval()

            accuracy_s, precision_s, recall_s, f1_s = [], [], [], []
            for i, (data, target) in enumerate(test_dataloader):
                data, target = data.to(decive), target.to(decive)

                pred = net(data)

                # 设置一个合理的阈值
                threshold = 0.5

                # 使用阈值进行二分类
                pred = (pred >= threshold).float()

                target, pred = target.detach().cpu(), pred.detach().cpu()

                accuracy, precision, recall, f1 = val.ret_calculate(target, pred)
                accuracy_s.append(accuracy)
                precision_s.append(precision)
                recall_s.append(recall)
                f1_s.append(f1)

            accuracy_p = torch.mean(torch.tensor(accuracy_s))
            precision_p = torch.mean(torch.tensor(precision_s))
            recall_p = torch.mean(torch.tensor(recall_s))
            f1_p = torch.mean(torch.tensor(f1_s))

            # 去掉第一组数据
            if epoch != 0:
                accuracy_list.append(accuracy_p)
                precision_list.append(precision_p)
                recall_list.append(recall_p)
                f1_list.append(f1_p)
            print("----accuracy = {}, ----precision = {}, ----recall = {}, ----f1 = {}".format(accuracy_p, precision_p,
                                                                                               recall_p, f1_p))
    return accuracy_list, precision_list, recall_list, f1_list


if __name__ == "__main__":
    filePath = "../../data/KC2.csv"
    accuracy_list, precision_list, recall_list, f1_list = nural_networks(filePath)
    plot_singleModel("Neural Networks", accuracy_list, precision_list, recall_list, f1_list)
    # 打印模型性能指标
    print("Accuracy:", accuracy_list)
    print("Precision:", precision_list)
    print("Recall:", recall_list)
    print("F1 Score:", f1_list)